//
// Created by Tony on 2024/5/26.
//

#include "SymbolFinder.h"
#include "../common/KeyWords.h"
#include "SymbolTable.h"
#include <cstdint>
#include <cstring>
#include <memory>
#include <vector>

namespace Compiler {

using std::make_shared;
using std::static_pointer_cast;

int32_t SymbolFinder::visit(Parser::Ast* ast) { return ast->accept(this); }

int32_t SymbolFinder::visit(Parser::BooleanLiteral* ast) {
    return mClassTable->getClassIdByName(KeyWords::BOOLEAN);
}

int32_t SymbolFinder::visit(Parser::CharLiteral* ast) {
    return mClassTable->getClassIdByName(KeyWords::CHAR);
}

int32_t SymbolFinder::visit(Parser::Flt64Literal* ast) {
    return mClassTable->getClassIdByName(KeyWords::DOUBLE);
}

int32_t SymbolFinder::visit(Parser::Int64Literal* ast) {
    return mClassTable->getClassIdByName(KeyWords::INT);
}

int32_t SymbolFinder::visit(Parser::NullLiteral* ast) {
    return mClassTable->getClassIdByName(KeyWords::NIL);
}

int32_t SymbolFinder::visit(Parser::StringLiteral* ast) {
    return mClassTable->getClassIdByName(KeyWords::STRING);
}

int32_t SymbolFinder::visit(Parser::Name* ast) {
    auto name = ast->name();
    auto location = mCurrentSymbolTable->where(name);
    if (location.index != -1) {
        auto symbol = mCurrentSymbolTable->get(location);
        ast->setSymbol(symbol);
        return symbol.classId;
    }

    setErrMsg("undefined identifier %s at %s", name.c_str(), ast->position().startStr().c_str());
    return -1;
}

int32_t SymbolFinder::visit(Parser::Operator* ast) { return 0; }

int32_t SymbolFinder::visit(Parser::BinaryExpr* ast) {
    const auto* op = ast->op()->op();

    if (strcmp(op, ".") == 0) {
        auto leftTypeId = ast->left()->accept(this);
        auto leftClass = mClassTable->getClassById(leftTypeId);
        shared_ptr<SymbolTable> leftClassSymbolTable;
        if (leftClass == mClass) {
            leftClassSymbolTable = mClassSymbolTable;
        } else {
            leftClassSymbolTable = mClassSymbolTable->getOuter()->getImport(leftClass->getName());
        }
        auto right = static_pointer_cast<Parser::Name>(ast->right());
        auto rightSymbolLoc = leftClassSymbolTable->where(right->name());
        auto rightSymbol = leftClassSymbolTable->get(rightSymbolLoc);
        right->setSymbol(rightSymbol);
        return rightSymbol.classId;
    }

    if (strcmp(op, "=") == 0) {
        ast->left()->accept(this);
        if (ast->left()->isBinaryExpr()) {
            auto left = static_pointer_cast<Parser::BinaryExpr>(ast->left());
            auto leftRight = static_pointer_cast<Parser::Name>(left->right());
            leftRight->setIsLeftNearbyAssign();
        } else {
            auto left = static_pointer_cast<Parser::Name>(ast->left());
            left->setIsLeftNearbyAssign();
        }
        return ast->right()->accept(this);
    }

    if (strcmp(op, "[") == 0) {
        ast->left()->accept(this);
        ast->right()->accept(this);
        return 0; // TODO return elementTypeId;
    }

    auto leftTypeId = ast->left()->accept(this);
    ast->right()->accept(this);
    return leftTypeId;
}

int32_t SymbolFinder::visit(Parser::TernaryExpr* ast) { return 0; }

int32_t SymbolFinder::visit(Parser::UnaryExpr* ast) {
    const auto* op = ast->op()->op();
    if (strcmp(op, KeyWords::NEW) == 0) {
        auto expr = static_pointer_cast<Parser::BinaryExpr>(ast->expr());
        expr->right()->accept(this);
        return expr->left()->accept(this);
    }

    return ast->expr()->accept(this);
}

int32_t SymbolFinder::visit(Parser::TypeTag* ast) {
    auto tag = ast->tag()->name();
    return mClassTable->getClassIdByName(tag);
}

int32_t SymbolFinder::visit(Parser::AccessFlags* ast) { return 0; }

int32_t SymbolFinder::visit(Parser::SimpleStmt* ast) {
    ast->getExpr()->accept(this);
    return 0;
}

int32_t SymbolFinder::visit(Parser::AssertStmt* ast) { return 0; }

int32_t SymbolFinder::visit(Parser::BlockStmt* ast) {
    auto offset = mCurrentSymbolTable->size();
    auto blockSymbolTable = make_shared<SymbolTable>(ScopeType::BLOCK, offset, mCurrentSymbolTable);
    mCurrentSymbolTable->addInner(blockSymbolTable);

    ast->setSymbolTable(blockSymbolTable);

    auto blockSymbolFinder =
        make_shared<SymbolFinder>(mClass, mClassTable, mClassSymbolTable, blockSymbolTable);

    for (auto& child : ast->getChildren()) {
        child->accept(blockSymbolFinder.get());
    }
    return 0;
}

int32_t SymbolFinder::visit(Parser::ForStmt* ast) {
    auto forSymTable = make_shared<SymbolTable>(ScopeType::BLOCK, mCurrentSymbolTable->size(),
                                                mCurrentSymbolTable);
    mCurrentSymbolTable->addInner(forSymTable);

    auto symbolFinder =
        make_shared<SymbolFinder>(mClass, mClassTable, mClassSymbolTable, forSymTable);

    for (auto& child : ast->getChildren()) {
        child->accept(symbolFinder.get());
    }
    return 0;
}

int32_t SymbolFinder::visit(Parser::IfStmt* ast) {
    for (auto& child : ast->getChildren()) {
        child->accept(this);
    }
    return 0;
}

int32_t SymbolFinder::visit(Parser::NullStmt* ast) { return 0; }

int32_t SymbolFinder::visit(Parser::VarDefStmt* ast) {
    auto varName = ast->getName();
    auto location = mCurrentSymbolTable->where(varName);
    if (location.index != -1) {
        setErrMsg("var <%s> redefined, at<%s>", ast->getName().c_str(),
                  ast->position().startStr().c_str());
        return -1;
    }

    int32_t typeId = 0;
    if (ast->getTypeTag() != nullptr) {
        auto className = ast->getTypeTag()->tag()->name();
        typeId = mClassTable->getClassIdByName(className);
    } else if (ast->getInitExpr() != nullptr) {
        typeId = ast->getInitExpr()->accept(this);
    } else {
        // not reachable
    }

    location = mCurrentSymbolTable->add(varName, SymbolType::LOCAL, typeId);
    ast->setSymbol(mCurrentSymbolTable->get(location));

    return typeId;
}

int32_t SymbolFinder::visit(Parser::ReturnStmt* ast) {
    auto returnExpr = ast->getReturnExpr();
    if (returnExpr != nullptr) {
        return returnExpr->accept(this);
    }
    return mClassTable->getClassIdByName(KeyWords::VOID);
}

int32_t SymbolFinder::visit(Parser::ArrayIndex* ast) { return 0; }

int32_t SymbolFinder::visit(Parser::ParameterList* ast) {
    vector<int32_t> types;
    for (int i = 0; i < (ast->size() >> 1); i++) {
        auto paramName = ast->name(i);
        auto location = mCurrentSymbolTable->where(paramName->name());
        if (location.index != -1) {
            setErrMsg("parameter <%s> redefined, at<%s>", paramName->name().c_str(),
                      paramName->position().startStr().c_str());
            return -1;
        }
        auto typeId = ast->typeTag(i)->accept(this);
        location = mCurrentSymbolTable->add(paramName->name(), SymbolType::LOCAL, typeId);
        paramName->setSymbol(mCurrentSymbolTable->get(location));
        types.push_back(typeId);
    }
    ast->setTypes(types);
    return 0;
}

int32_t SymbolFinder::visit(Parser::FunctionDefine* ast) {
    auto funcName = ast->name();

    auto location = mClassSymbolTable->where(funcName->name());
    if (location.index != -1) { // check if function has been defined
        setErrMsg("function %s redefined at %s", funcName->name().c_str(),
                  funcName->position().startStr().c_str());
        return -1;
    }

    auto funcSymbolTable = make_shared<SymbolTable>(ScopeType::METHOD, 0, mClassSymbolTable);
    auto funcSymbolFinder =
        make_shared<SymbolFinder>(mClass, mClassTable, mClassSymbolTable, funcSymbolTable);

    ast->parameterList()->accept(funcSymbolFinder.get());
    auto parameterTypes = ast->parameterList()->getTypes();

    auto returnType = ast->returnTypeTag()->accept(funcSymbolFinder.get());

    ast->block()->accept(funcSymbolFinder.get());

    auto function = make_shared<RtCode::Method>(funcName->name());
    function->addAccessFlag(RtCode::AccessFlag::STATIC);
    function->setType(parameterTypes, returnType);
    function->setArgSlotCount(parameterTypes.size());
    function->setLocalSlotCount(funcSymbolTable->maxSlotCount());

    mClass->addMethod(function);
    mClassSymbolTable->add(funcName->name(), SymbolType::METHOD, returnType);

    ast->setSymbolTable(funcSymbolTable);
    return returnType;
}

int32_t SymbolFinder::visit(Parser::ArgumentList* ast) {
    for (auto& child : ast->getChildren()) {
        child->accept(this);
    }
    return 0;
}

// module
int32_t SymbolFinder::visit(Parser::Module* ast) {
    for (auto& child : ast->getChildren()) {
        child->accept(this);
    }
    return 0;
}

int32_t SymbolFinder::visit(Parser::FieldDefine* ast) {
    auto fieldName = ast->getName();
    auto location = mClassSymbolTable->where(fieldName);
    if (location.index != -1) {
        setErrMsg("field <%s> redefined, at<%s>", ast->getName().c_str(),
                  ast->position().startStr().c_str());
        return -1;
    }

    int32_t typeId = 0;
    if (ast->getTypeTag() != nullptr) {
        auto className = ast->getTypeTag()->tag()->name();
        typeId = this->mClassTable->getClassIdByName(className);
    } else if (ast->getInitExpr() != nullptr) {
        typeId = ast->getInitExpr()->accept(this);
    } else {
        // not reachable
    }

    auto field = make_shared<RtCode::Field>(fieldName);
    field->setAccessFlag(ast->getAccessFlags()->getFlags());
    field->setOwnerClass(mClass);
    field->setType(typeId);

    mClass->addField(field);
    mClassSymbolTable->add(fieldName, SymbolType::FIELD, typeId);

    return typeId;
}

int32_t SymbolFinder::visit(Parser::MethodDefine* ast) {
    auto methodName = ast->getName();
    auto location = mClassSymbolTable->where(methodName);
    if (location.index != -1) {
        setErrMsg("method <%s> redefined, at<%s>", ast->getName().c_str(),
                  ast->position().startStr().c_str());
        return -1;
    }

    int32_t returnType;
    if (ast->getReturnTypeTag() != nullptr) {
        auto className = ast->getReturnTypeTag()->tag()->name();
        returnType = mClassTable->getClassIdByName(className);
    } else {
        returnType = mClassTable->getClassIdByName(KeyWords::VOID);
    }

    mClassSymbolTable->add(methodName, SymbolType::METHOD, returnType);
    auto method = make_shared<RtCode::Method>(methodName);
    method->setAccessFlag(ast->getAccessFlags()->getFlags());

    auto methodSymbolTable =
        make_shared<SymbolTable>(ScopeType::METHOD, 0, mClassSymbolTable, methodName);
    if (!method->isStatic()) {
        methodSymbolTable->add(KeyWords::THIS, SymbolType::LOCAL,
                               mClassTable->getClassIdByName(mClass->getName()));
    }
    auto methodSymbolFinder =
        make_shared<SymbolFinder>(mClass, mClassTable, mClassSymbolTable, methodSymbolTable);
    ast->getParameterList()->accept(methodSymbolFinder.get());
    ast->getBody()->accept(methodSymbolFinder.get());
    ast->setSymbolTable(methodSymbolTable);

    auto paramTypes = ast->getParameterList()->getTypes();
    method->setType(paramTypes, returnType);
    method->setArgSlotCount(method->isStatic() ? paramTypes.size() : paramTypes.size() + 1);
    method->setLocalSlotCount(methodSymbolTable->maxSlotCount());

    mClass->addMethod(method);
    mClassSymbolTable->add(methodName, SymbolType::METHOD, returnType);

    return returnType;
}

int32_t SymbolFinder::visit(Parser::NativeMethodDecl* ast) {
    auto methodName = ast->getName();
    auto location = mClassSymbolTable->where(methodName);
    if (location.index != -1) {
        setErrMsg("method <%s> redefined, at<%s>", ast->getName().c_str(),
                  ast->position().startStr().c_str());
        return -1;
    }

    int32_t returnType;
    if (ast->getReturnTypeTag() != nullptr) {
        auto className = ast->getReturnTypeTag()->tag()->name();
        returnType = mClassTable->getClassIdByName(className);
    } else {
        returnType = mClassTable->getClassIdByName(KeyWords::VOID);
    }

    location = mClassSymbolTable->add(methodName, SymbolType::METHOD, returnType);
    ast->setLocation(location);

    auto methodSymbolTable = make_shared<SymbolTable>(ScopeType::METHOD, 0, mClassSymbolTable);
    auto methodSymbolFinder =
        make_shared<SymbolFinder>(mClass, mClassTable, mClassSymbolTable, methodSymbolTable);
    ast->getParameterList()->accept(methodSymbolFinder.get());
    ast->setSymbolTable(methodSymbolTable);

    auto method = make_shared<RtCode::Method>(methodName);
    method->setAccessFlag(ast->getAccessFlags()->getFlags());
    auto paramTypes = ast->getParameterList()->getTypes();
    method->setType(paramTypes, returnType);
    auto argSlotCount = method->isStatic() ? paramTypes.size() : paramTypes.size() + 1;
    method->setArgSlotCount(argSlotCount);
    auto localSlotCount = method->isStatic() ? methodSymbolTable->maxSlotCount()
                                             : methodSymbolTable->maxSlotCount() + 1;
    method->setLocalSlotCount(localSlotCount);

    mClass->addMethod(method);
    mClassSymbolTable->add(methodName, SymbolType::METHOD, returnType);

    return returnType;
}

int32_t SymbolFinder::visit(Parser::ClassBody* ast) {
    for (auto& child : ast->getChildren()) {
        child->accept(this);
    }
    return 0;
}

int32_t SymbolFinder::visit(Parser::ClassDefine* ast) {
    mClass = make_shared<RtCode::Class>(ast->getName(), mClassTable);
    mClass->setAccessFlag(ast->getAccessFlags());
    auto classId = mClassTable->addClass(mClass);

    using Compiler::ScopeType;
    using Compiler::SymbolTable;
    mClassSymbolTable =
        make_shared<SymbolTable>(ScopeType::CLASS, 0, mImportSymbolTable, ast->getName());
    mClassSymbolTable->add(ast->getName(), SymbolType::CLASS, classId);

    auto clinitMethodSymbolTable = make_shared<SymbolTable>(ScopeType::METHOD, 0, mClassSymbolTable,
                                                            Runtime::CLASS_INIT_METHOD_NAME);
    auto initMethodSymbolTable = make_shared<SymbolTable>(ScopeType::METHOD, 0, mClassSymbolTable,
                                                          Runtime::OBJECT_INIT_METHOD_NAME);
    initMethodSymbolTable->add(KeyWords::THIS, SymbolType::LOCAL, classId);
    mClassSymbolTable->addInner(clinitMethodSymbolTable);
    mClassSymbolTable->addInner(initMethodSymbolTable);

    ast->setClassId(classId);
    ast->getClassBody()->accept(this);

    auto method = make_shared<RtCode::Method>(Runtime::CLASS_INIT_METHOD_NAME);
    method->setLocalSlotCount(clinitMethodSymbolTable->maxSlotCount());
    mClass->addMethod(method);

    method = make_shared<RtCode::Method>(Runtime::OBJECT_INIT_METHOD_NAME);
    method->setLocalSlotCount(initMethodSymbolTable->maxSlotCount());
    method->setArgSlotCount(1); // only arg is this
    mClass->addMethod(method);

    return classId;
}

int32_t SymbolFinder::visit(Parser::ClassModule* ast) {
    mImportSymbolTable = std::make_shared<SymbolTable>(ScopeType::IMPORTS, 0, nullptr);
    for (auto& import : ast->getImports()) {
        auto className = import->getClassName();
        auto classId = mClassTable->getClassIdByName(className);
        auto clazz = mClassTable->getClassById(classId);
        mImportSymbolTable->add(className, SymbolType::CLASS, classId);
        mImportSymbolTable->addImport(clazz->getSymbolTable());
    }
    for (auto& classDefine : ast->getClassDefines()) {
        classDefine->accept(this);
    }
    return 0;
}

int32_t SymbolFinder::visit(Parser::ViewBuilder* ast) { return 0; }
int32_t SymbolFinder::visit(Parser::ViewBuilderList* ast) { return 0; }
int32_t SymbolFinder::visit(Parser::ViewBuilderBlock* ast) { return 0; }
int32_t SymbolFinder::visit(Parser::AttributeSetter* ast) { return 0; }
int32_t SymbolFinder::visit(Parser::AttributeSetterList* ast) { return 0; }

} // namespace Compiler
